//https://www.gnu.org/software/gsl/doc/html/spmatrix.html

#include "toroidal.h"
#include <stdio.h>
#include "math.h"
#include <gsl/gsl_spmatrix.h>
#include <gsl/gsl_spblas.h>
#include <gsl/gsl_vector.h>

gsl_spmatrix* partr()
{
    double ivdr = 1.0/dr;
    int mdim = Nx*Ny;
    
    gsl_spmatrix *pg = gsl_spmatrix_alloc(2*mdim, 2*mdim);
    for (int i = 0; i < Nx*Ny; i++)
    {
        int ix = i/Ny;
        int lf = (ix == 0 ? i : i-Ny);
        int rg = (ix == Nx-1 ? i : i+Ny);
        double rgap = 1.0/(double)((rg - lf)/Ny);

        gsl_spmatrix_set(pg, i, lf, -ivdr * rgap);
        gsl_spmatrix_set(pg, i, rg, ivdr * rgap);
        gsl_spmatrix_set(pg, i+mdim, lf+mdim, -ivdr * rgap);
        gsl_spmatrix_set(pg, i+mdim, rg+mdim, ivdr * rgap);
    }
    return pg;
}

gsl_spmatrix* partz()
{
    double ivdz = 1.0/dz;
    int mdim = Nx*Ny;
    gsl_spmatrix *pg = gsl_spmatrix_alloc(2*mdim, 2*mdim);
    for (int i = 0; i < Nx*Ny; i++)
    {
        int iy = i%Ny;
        int dn = (iy == 0 ? i : i-1);
        int up = (iy == Ny-1 ? i : i+1);
        double zgap = 1.0/(double)(up-dn);

        gsl_spmatrix_set(pg, i, dn, -ivdz * zgap);
        gsl_spmatrix_set(pg, i, up, ivdz * zgap);
        gsl_spmatrix_set(pg, i+mdim, dn+mdim, -ivdz * zgap);
        gsl_spmatrix_set(pg, i+mdim, up+mdim, ivdz * zgap);
    }
    return pg;
}

gsl_spmatrix* paragrad(vfield b, int n)
{
    double* pbr = b.getfr();
    double* pbp = b.getfphi();
    double* pbz = b.getfz();
    double ivdr = 1.0/dr;
    double ivdz = 1.0/dz;
    int mdim = Nx*Ny;
    
    gsl_spmatrix *pg = gsl_spmatrix_alloc(2*mdim, 2*mdim);
    for (int i = 0; i < Nx*Ny; i++)
    {
        int ix = i/Ny;
        int iy = i%Ny;
        int dn = (iy == 0 ? i : i-1);
        int up = (iy == Ny-1 ? i : i+1);
        int lf = (ix == 0 ? i : i-Ny);
        int rg = (ix == Nx-1 ? i : i+Ny);
        double rgap = 1.0/(double)((rg - lf)/Ny);
        double zgap = 1.0/(double)(up-dn);
        double R = (i/Ny)*dr + Rmin;

        gsl_spmatrix_set(pg, i, dn, -pbz[i] *ivdz * zgap);
        gsl_spmatrix_set(pg, i, up, pbz[i] *ivdz * zgap);
        gsl_spmatrix_set(pg, i, lf, -pbr[i] * ivdr * rgap);
        gsl_spmatrix_set(pg, i, rg, pbr[i] * ivdr * rgap);

        gsl_spmatrix_set(pg, i+mdim, dn+mdim, -pbz[i] * ivdz * zgap);
        gsl_spmatrix_set(pg, i+mdim, up+mdim, pbz[i] * ivdz * zgap);
        gsl_spmatrix_set(pg, i+mdim, lf+mdim, -pbr[i] * ivdr * rgap);
        gsl_spmatrix_set(pg, i+mdim, rg+mdim, pbr[i] * ivdr * rgap);

        gsl_spmatrix_set(pg, i, i+mdim, pbp[i]*n/R);
        gsl_spmatrix_set(pg, i+mdim, i, -pbp[i]*n/R);
    }
    return pg;
}

gsl_spmatrix* laplacian(int mdim)
{
    gsl_spmatrix* pr, *pz, *prr, *pzz;
    prr = gsl_spmatrix_ccs(gsl_spmatrix_alloc(2*mdim, 2*mdim));
    pzz = gsl_spmatrix_ccs(gsl_spmatrix_alloc(2*mdim, 2*mdim));
    pr = gsl_spmatrix_ccs(partr());
    pz = gsl_spmatrix_ccs(partz());

    gsl_spblas_dgemm(1.0, pr,pr,prr);
    gsl_spblas_dgemm(1.0, pz,pz,pzz);
    
    gsl_spmatrix* lap = gsl_spmatrix_alloc(2*mdim, 2*mdim);
    lap = gsl_spmatrix_ccs(lap);
    gsl_spmatrix_add(lap, prr, pzz);
    gsl_spmatrix_free(pzz);
    gsl_spmatrix_free(prr);
    gsl_spmatrix_free(pr);
    gsl_spmatrix_free(pz);
    return lap;
}

gsl_spmatrix* diag(gsl_vector* a, int mdim)
{
    gsl_spmatrix *pg = gsl_spmatrix_alloc(2*mdim, 2*mdim);
    for (int i = 0; i < mdim; i++)
    {
        gsl_spmatrix_set(pg, i, i, a->data[i]);
        gsl_spmatrix_set(pg, i+mdim, i+mdim, a->data[i+mdim]);
    }
    return pg;
}

void world::init()
{
    this->partialr = partr();
    this->partialz = partz();

}

void world::shear(fiq phiiq, vfiq B1iq)
{
    gsl_spmatrix* A;
    A = gsl_spmatrix_ccs(paragrad(*(this->b), phiiq.n));
    int mdim = Nx*Ny;
    double* pxi = phiiq.get_fi();
    double* pxq = phiiq.get_fq();
    gsl_vector *x =  gsl_vector_alloc(2*mdim);
    gsl_vector *y =  gsl_vector_alloc(2*mdim);
    gsl_vector *u =  gsl_vector_alloc(2*mdim);
    gsl_vector *v =  gsl_vector_alloc(2*mdim);
    for(int i = 0; i < mdim; i++)
    {
        gsl_vector_set(x, i, pxi[i]);
        gsl_vector_set(x, i+mdim, pxq[i]);
    }
    gsl_spmatrix *pr = gsl_spmatrix_ccs(partialr);
    gsl_spmatrix *pz = gsl_spmatrix_ccs(partialz);
    gsl_spblas_dgemv(CblasNoTrans,1.0,A,x, 0.0,y);
    gsl_spblas_dgemv(CblasNoTrans,1.0,pr,y,0.0,u);
    gsl_spblas_dgemv(CblasNoTrans,1.0,pz,y,0.0,v);
    
    double* B1ri = B1iq.get_ifr();
    double* B1zi = B1iq.get_ifz();
    double* B1rq = B1iq.get_qfr();
    double* B1zq = B1iq.get_qfz();
    for(int i = 0; i < mdim; i++)
    {   
        B1ri[i] = gsl_vector_get(u, i);
        B1rq[i] = gsl_vector_get(u, i+mdim);
        B1zi[i] = gsl_vector_get(v, i);
        B1zq[i] = gsl_vector_get(v, i+mdim);
    }
    return;
}

void world::flow(fiq phiiq, vfiq viq)
{
    int mdim = Nx*Ny;
    double* pB = this->Bsca->getdata();
    double* pxi = phiiq.get_fi();
    double* pxq = phiiq.get_fq();
    gsl_vector *x =  gsl_vector_alloc(2*mdim);
    gsl_vector *u =  gsl_vector_alloc(2*mdim);
    gsl_vector *v =  gsl_vector_alloc(2*mdim);
    gsl_vector *ivBsc = gsl_vector_alloc(2*mdim);
    for(int i = 0; i < mdim; i++)
    {
        gsl_vector_set(x, i, pxi[i]);
        gsl_vector_set(x, i+mdim, pxq[i]);
        if(pB[i] != 0.0)
        {
            gsl_vector_set(ivBsc, i, 1.0/pB[i]);
            gsl_vector_set(ivBsc, i+mdim, 1.0/pB[i]);
        } else
        {
            gsl_vector_set(ivBsc, i, 0.0);
            gsl_vector_set(ivBsc, i+mdim, 0.0);
        }
    }
    gsl_spmatrix *pr = gsl_spmatrix_ccs(partialr);
    gsl_spmatrix *pz = gsl_spmatrix_ccs(partialz);
    gsl_spmatrix* ivBst = gsl_spmatrix_ccs(diag(ivBsc, mdim));
    gsl_spmatrix* op1 = gsl_spmatrix_ccs(gsl_spmatrix_alloc(2*mdim, 2*mdim));
    gsl_spmatrix* op2 = gsl_spmatrix_ccs(gsl_spmatrix_alloc(2*mdim, 2*mdim));
    gsl_spblas_dgemm(1.0, ivBst, pr, op1);
    gsl_spblas_dgemm(1.0, ivBst, pz, op2);
    gsl_spblas_dgemv(CblasNoTrans,1.0,op1,x,0.0,u);
    gsl_spblas_dgemv(CblasNoTrans,1.0,op2,x,0.0,v);
    
    double* B1ri = viq.get_ifr();
    double* B1zi = viq.get_ifz();
    double* B1rq = viq.get_qfr();
    double* B1zq = viq.get_qfz();
    for(int i = 0; i < mdim; i++)
    {   
        B1ri[i] = gsl_vector_get(u, i);
        B1rq[i] = gsl_vector_get(u, i+mdim);
        B1zi[i] = gsl_vector_get(v, i);
        B1zq[i] = gsl_vector_get(v, i+mdim);
    }
    return;
}

void perpgrad(fiq xiq, vfiq Eiq)
{
    int mdim = Nx*Ny;

    double* pxi = xiq.get_fi();
    double* pxq = xiq.get_fq();
    gsl_vector *x =  gsl_vector_alloc(2*mdim);
    gsl_vector *u =  gsl_vector_alloc(2*mdim);
    gsl_vector *v =  gsl_vector_alloc(2*mdim);
    for(int i = 0; i < mdim; i++)
    {
        gsl_vector_set(x, i, pxi[i]);
        gsl_vector_set(x, i+mdim, pxq[i]);
    }
    gsl_spmatrix *pr = gsl_spmatrix_ccs(partr());
    gsl_spmatrix *pz = gsl_spmatrix_ccs(partz());
    gsl_spblas_dgemv(CblasNoTrans,1.0,pr,x,0.0,u);
    gsl_spblas_dgemv(CblasNoTrans,1.0,pz,x,0.0,v);
    
    double* Eri = Eiq.get_ifr();
    double* Ezi = Eiq.get_ifz();
    double* Erq = Eiq.get_qfr();
    double* Ezq = Eiq.get_qfz();
    for(int i = 0; i < mdim; i++)
    {   
        Eri[i] = gsl_vector_get(u, i);
        Erq[i] = gsl_vector_get(u, i+mdim);
        Ezi[i] = gsl_vector_get(v, i);
        Ezq[i] = gsl_vector_get(v, i+mdim);
    }
    return;
}

void world::flb_gsl(fiq xiq, fiq yiq) //field line bending term
{
    int mdim = Nx*Ny;
    double* pxi = xiq.get_fi();
    double* pxq = xiq.get_fq();
    double* pB = this->Bsca->getdata();
    gsl_vector *x =  gsl_vector_alloc(2*mdim);
    gsl_vector *y =  gsl_vector_alloc(2*mdim);
    gsl_vector *ivBsc = gsl_vector_alloc(2*mdim);
    for(int i = 0; i < mdim; i++)
    {
        gsl_vector_set(x, i, pxi[i]);
        gsl_vector_set(x, i+mdim, pxq[i]);
        if(pB[i] != 0.0)
        {
            gsl_vector_set(ivBsc, i, 1.0/pB[i]);
            gsl_vector_set(ivBsc, i+mdim, 1.0/pB[i]);
        } else
        {
            gsl_vector_set(ivBsc, i, 0.0);
            gsl_vector_set(ivBsc, i+mdim, 0.0);
        }
    }

    gsl_spmatrix  *A, *lap, *flbsp, *ivBst, *C, *D;
    lap = gsl_spmatrix_alloc(2*mdim, 2*mdim);
    A = gsl_spmatrix_ccs(paragrad(*(this->b), xiq.n));
    ivBst = gsl_spmatrix_ccs(diag(ivBsc, mdim));
    C = gsl_spmatrix_ccs(gsl_spmatrix_alloc(2*mdim, 2*mdim));
    D = gsl_spmatrix_ccs(gsl_spmatrix_alloc(2*mdim, 2*mdim));
    flbsp = gsl_spmatrix_ccs(gsl_spmatrix_alloc(2*mdim, 2*mdim));
    lap = laplacian(mdim);
    
    gsl_spblas_dgemm(1.0, lap, A, C);
    gsl_spblas_dgemm(1.0, ivBst, C, D);
    gsl_spblas_dgemm(1.0, A, C, flbsp);
    gsl_spblas_dgemv(CblasNoTrans,1.0,flbsp,x,0.0,y);

    double* pyi = yiq.get_fi();
    double* pyq = yiq.get_fq();
    
    //double sum = 0.0;
    for(int i = 0; i < mdim; i++)
    {   
        //sum += abs(gsl_vector_get(y, i));
        pyi[i] = gsl_vector_get(y, i);
        pyq[i] = gsl_vector_get(y, i+mdim);
    }
    gsl_spmatrix_free(A);
    gsl_spmatrix_free(C);
    gsl_spmatrix_free(lap);
    gsl_spmatrix_free(flbsp);
    return;
}

void world::inertial_gsl(fiq xiq, fiq yiq)
{
    int mdim = Nx*Ny;
    double* pxi = xiq.get_fi();
    double* pxq = xiq.get_fq();
    double* pva = this->vA->getdata();
    gsl_vector *x =  gsl_vector_alloc(2*mdim);
    gsl_vector *y =  gsl_vector_alloc(2*mdim);
    gsl_vector *va = gsl_vector_alloc(2*mdim);
    for(int i = 0; i < mdim; i++)
    {
        gsl_vector_set(x, i, pxi[i]);
        gsl_vector_set(x, i+mdim, pxq[i]);
        if(pva[i] == 0.0)
        {
            gsl_vector_set(va, i, 0.0);
            gsl_vector_set(va, i+mdim, 0.0);
        }
        else
        {
            gsl_vector_set(va, i, 1.0/pva[i]);
            gsl_vector_set(va, i+mdim, 1.0/pva[i]);
        }
    }

    gsl_spmatrix* pr, *pz, *prr, *pzz, *invva, *temp1, *temp2;
    prr = gsl_spmatrix_ccs(gsl_spmatrix_alloc(2*mdim, 2*mdim));
    pzz = gsl_spmatrix_ccs(gsl_spmatrix_alloc(2*mdim, 2*mdim));
    temp1 = gsl_spmatrix_ccs(gsl_spmatrix_alloc(2*mdim, 2*mdim));
    temp2 = gsl_spmatrix_ccs(gsl_spmatrix_alloc(2*mdim, 2*mdim));
    pr = gsl_spmatrix_ccs(partr());
    pz = gsl_spmatrix_ccs(partz());
    invva = gsl_spmatrix_ccs(diag(va, mdim));

    gsl_spblas_dgemm(1.0, invva, pr, temp1);
    gsl_spblas_dgemm(1.0, pr, temp1, prr);
    gsl_spblas_dgemm(1.0, invva, pz, temp1);
    gsl_spblas_dgemm(1.0, pz, temp1, pzz);
    
    gsl_spmatrix* iner = gsl_spmatrix_alloc(2*mdim, 2*mdim);
    iner = gsl_spmatrix_ccs(iner);
    gsl_spmatrix_add(iner, prr, pzz);
    
    gsl_spblas_dgemv(CblasNoTrans,1.0,iner,x,0.0,y);
    
    double* pyi = yiq.get_fi();
    double* pyq = yiq.get_fq();
    
    for(int i = 0; i < mdim; i++)
    {   
        pyi[i] = gsl_vector_get(y, i);
        pyq[i] = gsl_vector_get(y, i+mdim);
    }
    gsl_spmatrix_free(pr);
    gsl_spmatrix_free(pz);
    gsl_spmatrix_free(prr);
    gsl_spmatrix_free(pzz);
    gsl_spmatrix_free(temp1);
    gsl_spmatrix_free(temp2);
    return;
}

void world::partr_noiq(double* px, double* py)
{
    int mdim = Nx*Ny;
    double ivdr = 1.0/dr;
    gsl_spmatrix *pg = gsl_spmatrix_alloc(mdim, mdim);
    for (int i = 0; i < Nx*Ny; i++)
    {
        int ix = i/Ny;
        int lf = (ix == 0 ? i : i-Ny);
        int rg = (ix == Nx-1 ? i : i+Ny);
        double rgap = 1.0/(double)((rg - lf)/Ny);
        gsl_spmatrix_set(pg, i, lf, -ivdr * rgap);
        gsl_spmatrix_set(pg, i, rg, ivdr * rgap);
    }
    gsl_vector *x =  gsl_vector_alloc(mdim);
    gsl_vector *y =  gsl_vector_alloc(mdim);
    for(int i = 0; i < mdim; i++)
        gsl_vector_set(x, i, px[i]);
        
    gsl_spblas_dgemv(CblasNoTrans,1.0,pg, x, 0.0,y);
    for(int i = 0; i < mdim; i++)
        py[i] = gsl_vector_get(y, i);
    return;
}

void world::partz_noiq(double* px, double* py)
{
    double ivdz = 1.0/dz;
    int mdim = Nx*Ny;
    gsl_spmatrix *pg = gsl_spmatrix_alloc(mdim, mdim);
    for (int i = 0; i < Nx*Ny; i++)
    {
        int iy = i%Ny;
        int dn = (iy == 0 ? i : i-1);
        int up = (iy == Ny-1 ? i : i+1);
        double zgap = 1.0/(double)(up-dn);
        gsl_spmatrix_set(pg, i, dn, -ivdz * zgap);
        gsl_spmatrix_set(pg, i, up, ivdz * zgap);
    }
    gsl_vector *x =  gsl_vector_alloc(mdim);
    gsl_vector *y =  gsl_vector_alloc(mdim);
    for(int i = 0; i < mdim; i++)
        gsl_vector_set(x, i, px[i]);
        
    gsl_spblas_dgemv(CblasNoTrans,1.0,pg, x, 0.0,y);
    for(int i = 0; i < mdim; i++)
        py[i] = gsl_vector_get(y, i);
    return;
}
